"""
存放关于图片识别方法
"""
import os

import cv2
import numpy as np
import pytesseract

from libs import work_spec


class Cv2Script:

    def __init__(self, driver):
        self.dr = driver

    def current_screen(self):
        r = self.dr.get_screenshot_as_png()
        print("屏幕信息获取完成！")
        img_array = np.fromstring(r, np.uint8)
        return cv2.imdecode(img_array, cv2.IMREAD_COLOR)

    def find_target_in_current_screen(self, temp_pic_name, threshold=0.95, debug_pic=None):
        """接触封印"""
        re_location = []
        template_img = cv2.imread(os.path.join(work_spec, temp_pic_name))
        if debug_pic:
            target_img = cv2.imread(os.path.join(work_spec, debug_pic))
        else:
            target_img = self.current_screen()

        # 获取目标图像和模板图像的宽高
        template_h, template_w = template_img.shape[:2]

        # 使用平方差匹配算法，也可以转成灰度图像进行比对
        result = cv2.matchTemplate(template_img, target_img, cv2.TM_CCOEFF_NORMED)

        # 设置匹配阈值
        threshold = threshold

        # 找到匹配结果中超过阈值的位置
        locations = np.where(result >= threshold)
        print("当前坐标元素:", locations)
        locations = list(zip(*locations[::-1]))
        # 在目标图像上绘制矩形框
        for loc in locations:
            top_left = loc
            re_location.append((int(top_left[0] + template_w / 2), int(top_left[1] + template_h / 2)))
            if debug_pic:
                bottom_right = (top_left[0] + template_w, top_left[1] + template_h)
                cv2.rectangle(target_img, top_left, bottom_right, (0, 255, 0), 2)
        if debug_pic:
            # 显示结果图像
            cv2.imshow('Result', target_img)
            cv2.waitKey(0)
            cv2.destroyAllWindows()
        return re_location

    def get_pic_str(self, left, upper, right, lower):
        screen = self.current_screen()
        cropped = screen[upper:lower, left:right]
        # cropped = Image.open(os.path.join(work_spec, "jcfy.png"))
        tessdata_dir_config = '--tessdata-dir "C:\\Program Files\\Tesseract-OCR\\tessdata"'
        text = pytesseract.image_to_string(cropped, config=tessdata_dir_config)
        print(text)
        return text
